import java.util.Random;

/**
 * 528. 按权重随机选择
 * https://leetcode-cn.com/problems/random-pick-with-weight/
 */
public class Solutions_528 {
    public static void main(String[] args) {
//        int[] arr = {1, 3};
//        int[] arr = {1};
        int[] arr = {3, 14, 1, 7};

        Solution_528 solution = new Solution_528(arr);
        int res1 = solution.pickIndex();  // output:1，返回下标 1，返回该下标概率为 3/4
        System.out.println(res1);
        int res2 = solution.pickIndex();  // output: 1
        System.out.println(res2);
        int res3 = solution.pickIndex();  // output: 1
        System.out.println(res3);
        int res4 = solution.pickIndex();  // output: 1
        System.out.println(res4);
        int res5 = solution.pickIndex();  // output: 0，返回下标 0，返回该下标概率为 1/4
        System.out.println(res5);
    }
}

/**
 * 解题思路：前缀和 + 二分查找
 */
class Solution_528 {
    int[] preSum;
    Random random;
    int sum;

    public Solution_528(int[] w) {
        this.preSum = new int[w.length];
        preSum[0] = w[0];
        for (int i = 1; i < w.length; i++) {
            preSum[i] = w[i] + preSum[i - 1];
        }
        random = new Random();
        sum = preSum[w.length - 1];
    }

    public int pickIndex() {
        // 随机数区间为：[1, sum]
        int num = random.nextInt(sum) + 1;
//        System.out.println("num:" + num);
        // 二分查找
        int left = 0, right = preSum.length - 1;
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (preSum[mid] < num) {
                left = mid + 1;
            } else {
                right --;
            }
        }
        return left;
    }
}
